{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from lmfit import minimize, Parameters, Parameter, report_fit\n",
    "from scipy.integrate import odeint\n",
    "import sys\n",
    "import time\n",
    "import datetime\n",
    "from datetime import timedelta  \n",
    "from mpi4py import MPI # use the Discovery cluster @ Dartmouth to perform parallel processing\n",
    "import import_ipynb\n",
    "from utils_common import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "_Figure_PATH_XX_ = './figures/parameters_xx/'\n",
    "_Data_PATH_XX_ = './data/parameters_xx/'\n",
    "\n",
    "_Network_P2P_PATH_ = './data/data_network_P2P.csv'\n",
    "_City_PATH_ = './data/data_DXY_city_all.csv'\n",
    "_Province_PATH_ = './data/data_DXY_province_all.csv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_city, data_province, data_province_domestic = load_DXY_raw()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SEIR simulation WITHOUT MIGRATION \n",
    "# same timeline for all porvinces\n",
    "# beta, gamma and sigma are time-dependent\n",
    "# beta: first changes every TI = 3 days for T = 21 days, then changes every ti days\n",
    "# gamma and sigma: first changes every TT = 6 days for 2*TT = 12 days, then stays the same\n",
    "\n",
    "# possible algorithms:\n",
    "# ’leastsq’: Levenberg-Marquardt (default)\n",
    "# ’emcee’: Maximum likelihood via Monte-Carlo Markov Chain\n",
    "# ’dual_annealing’: Dual Annealing optimization\n",
    "\n",
    "def simulation_province_xx(df = data_province_domestic, update_date_tr = False, rank = 0, # the index of the province\n",
    "                           T = 21, TT = 6, TI = 3, ti = 7, mmethod = 'dual_annealing', kws = False):\n",
    "    \n",
    "    start_time = time.time()\n",
    "    province = names_province[rank]\n",
    "    \n",
    "    df = df[df.update_date <= update_date_tr]\n",
    "    df_single = df[df.province_name_en == province]\n",
    "    \n",
    "    timespan_local = (max(df_single.update_date) - min(df_single.update_date)).days + 1 # provincial timespan\n",
    "    timespan = (max(df.update_date) - min(df.update_date)).days + 1 # national timespan\n",
    "    \n",
    "    list_front = np.repeat(range(int(np.ceil(T/TI))), TI).tolist()\n",
    "    list_end = np.repeat(range(int(np.ceil((timespan - T)/ti))), ti).tolist()\n",
    "    list_end = [item + max(list_front) + 1 for item in list_end]\n",
    "    index_list = list_front + list_end\n",
    "    index_list = index_list[:timespan] # index list for beta\n",
    "    \n",
    "    list_front_gs = [0]*TT + [1]*TT\n",
    "    list_end_gs = [2]*(timespan - TT - TT)\n",
    "    index_list_gs = list_front_gs + list_end_gs # index list for gamma and sigma\n",
    "    \n",
    "    # Least-square fitting\n",
    "    # SEIR model without migration\n",
    "    def SEIR_equations(states, t, ps):\n",
    "        \"\"\"Receptor synthesis-internalization model.\"\"\"\n",
    "        # integrate.ode may generate t values outside the data range\n",
    "        beta_index = index_list[min(round(t), timespan - 1)]\n",
    "        gs_index = index_list_gs[min(round(t), timespan - 1)]\n",
    "    \n",
    "        beta = ps['beta_' + str(beta_index)].value\n",
    "        gamma = ps['gamma_' + str(gs_index)].value\n",
    "        sigma = ps['sigma_' + str(gs_index)].value\n",
    "        N = ps['N'].value\n",
    "        S, E, I, R = states\n",
    "\n",
    "        return [-beta*S*I/N, beta*S*I/N - sigma*E, sigma*E - gamma*I, gamma*I]\n",
    "    \n",
    "    def SEIR(states_0, ts, ps):\n",
    "        \"\"\"\n",
    "        Solution to the ODE x'(t) = f(t,x,k) with initial condition x(0) = x0\n",
    "        \"\"\"\n",
    "        states = odeint(SEIR_equations, states_0, ts, args=(ps,))\n",
    "        return states\n",
    "    \n",
    "    # Calculate the residual\n",
    "    def residual(ps, ts, data):\n",
    "        states_0 = ps['S_0'].value, ps['E_0'].value, ps['I_0'].value, ps['R_0'].value\n",
    "        model = SEIR(states_0, ts, ps)\n",
    "        est = model[:,3]\n",
    "        est = est[-timespan_local:]\n",
    "        return (est - data).ravel()\n",
    "    \n",
    "    # Calculate the estimation\n",
    "    def estimation(ps, ts):\n",
    "        states_0 = ps['S_0'].value, ps['E_0'].value, ps['I_0'].value, ps['R_0'].value\n",
    "        model = SEIR(states_0, ts, ps)\n",
    "        est = model[:,3]\n",
    "        return est\n",
    "\n",
    "    ts = range(0, timespan)\n",
    "    \n",
    "    nbeta = int(T/TI + np.ceil((timespan - T)/ti)) # number of beta's\n",
    "    \n",
    "    data = np.array(df_single.cum_confirmed)\n",
    "    \n",
    "    # set parameters (incluing their bounds)\n",
    "    ps_bound = Parameters()\n",
    "    ps_bound.add(name = 'N', value = provincial_population_dict.get(province), vary = False)\n",
    "    ps_bound.add(name = 'E_0', value = 50, min = 0, max = 500)\n",
    "    ps_bound.add(name = 'I_0', value = 50, min = 0, max = 200)\n",
    "    ps_bound.add(name = 'R_0', value = 0, min = 0, max = 100)\n",
    "    ps_bound.add(name = 'S_0', expr = 'N - E_0 - I_0 - R_0')\n",
    "    for i in range(max(index_list) + 1):\n",
    "        ps_bound.add(name = 'beta_' + str(int(i)) , value = 0.5, min = 0.01, max = 1)\n",
    "    for i in range(max(index_list_gs) + 1):\n",
    "        ps_bound.add(name = 'gamma_' + str(int(i)), value = 0.5, min = 0.05, max = 0.5)\n",
    "    for i in range(max(index_list_gs) + 1):\n",
    "        ps_bound.add(name = 'sigma_' + str(int(i)), value = 0.5, min = 0.05, max = 0.5)\n",
    "    \n",
    "    # fit the model and find estimated values\n",
    "    if kws == False: # whether any key word is assigned to the algorithm\n",
    "        result = minimize(residual, ps_bound, args = (ts, data), method = mmethod)\n",
    "    elif mmethod == 'emcee':\n",
    "        result = minimize(residual, ps_bound, args = (ts, data), method = mmethod, steps = kws.get('steps'))\n",
    "    elif mmethod == 'dual_annealing':\n",
    "        result = minimize(residual, ps_bound, args = (ts, data), method = mmethod, maxiter = kws.get('maxiter'), initial_temp = kws.get('initial_temp'))\n",
    "    elif mmethod == 'differential_evolution':\n",
    "        result = minimize(residual, ps_bound, args = (ts, data), method = mmethod, maxiter = kws.get('maxiter'))\n",
    "        \n",
    "    residual = result.residual.reshape(data.shape)\n",
    "    final = estimation(result.params, ts)\n",
    "    data = np.concatenate(([None]*(timespan - timespan_local), data), axis=0)\n",
    "    residual = np.concatenate(([None]*(timespan - timespan_local), residual), axis=0)\n",
    "\n",
    "    # plot real data and fitted curves\n",
    "    palette = plt.get_cmap('magma')\n",
    "    fig = plt.figure(figsize = (10,5))\n",
    "    plt.plot(ts, data, '-o', color = palette(0.6))\n",
    "    plt.plot(ts, final, linewidth = 2, color = palette(0.5))\n",
    "    plt.title(province)\n",
    "    plt.xlabel('Date')\n",
    "    plt.ylabel('Number of people')\n",
    "    fig.savefig(_Figure_PATH_XX_ + province + '_conf_SEIR.pdf', dpi = 400) # save the figure\n",
    "    \n",
    "    print(\"processor %d of %d\" % (rank, size))\n",
    "    print(\"--- %s seconds ---\" % (time.time() - start_time))\n",
    "    \n",
    "    df_parameters = pd.DataFrame(dict(result.params.valuesdict()).items(), columns=['parameter', 'value'])\n",
    "    df_parameters.to_csv(_Data_PATH_XX_ + province + '_parameters.csv', index = False) # save the parameters\n",
    "    df_final = pd.DataFrame({'estimation': final, 'real': data, 'residual': residual}, columns=['estimation', 'real', 'residual'])\n",
    "    df_final.to_csv(_Data_PATH_XX_ + province + '_estimation.csv', index = False) # save the estimated values\n",
    "    return result, df_parameters, df_final"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "#emcee_kws = dict(steps = 10)\n",
    "#differential_evolution_kws = dict(maxiter = 1)\n",
    "annealing_kws = dict(maxiter = 1, initial_temp = 1e4)\n",
    "date_tr = datetime.date(int(2020),int(3),int(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for rank, name in enumerate(names_province):\n",
    "    #result, df_parameters, df_final = simulation_province_xx(df = data_province_domestic, update_date_tr = date_tr, rank = rank, # the index of the province\n",
    "                           #T = 21, TT = 6, TI = 3, ti = 7, mmethod = 'dual_annealing', kws = annealing_kws)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parallel computing using MPI\n",
    "# settings are given in the pbs file\n",
    "comm = MPI.COMM_WORLD\n",
    "rank = comm.Get_rank()\n",
    "size = comm.Get_size()\n",
    "if size != len(names_province):\n",
    "    print(\"Error!\")\n",
    "    sys.exit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pbs file\n",
    "#!/bin/bash -l\n",
    "#PBS -A Fulab\n",
    "#PBS -q default\n",
    "#PBS -N worker29\n",
    "#PBS -l nodes=2:ppn=16\n",
    "#PBS -l walltime=480:00:00\n",
    "#module load python/3-Anaconda\n",
    "#source activate axel\n",
    "#cd $PBS_O_WORKDIR\n",
    "#mpirun -np 30 ./utils_parameters.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "result, df_parameters, df_final = simulation_province_xx(df = data_province_domestic, update_date_tr = date_tr, rank = rank, \n",
    "                                                        T = 21, TT = 6, TI = 3, ti = 7, mmethod = 'dual_annealing', kws = annealing_kws)\n",
    "#print(\"--- %s seconds ---\" % (time.time() - start))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
